from langchain_huggingface import HuggingFaceEmbeddings  # 替换原langchain_community的导入
from langchain_chroma import Chroma
from typing import List, Dict, Any
import os

class ChromaDBManager:
    def __init__(self, persist_dir: str = "./chroma_db", collection_name: str = "analysis_reports"):
        self.embedding = HuggingFaceEmbeddings(
            model_name="BAAI/bge-large-zh-v1.5",
            model_kwargs={'device': 'cpu'},
            # encode_kwargs={'normalize_embeddings': True}
        )
        os.makedirs(persist_dir, exist_ok=True)
        self.vector_db = Chroma(
            persist_directory=persist_dir,
            embedding_function=self.embedding,
            collection_metadata={"hnsw:space": "cosine"},  # ✅ 新增此行
            collection_name=collection_name
        )
    def store_data(self, documents: List[str], metadatas: List[Dict[str, Any]] = None):
        """存储数据到数据库"""
        if metadatas and len(documents) != len(metadatas):
            raise ValueError("Documents and metadatas must have the same length")
        
        # 添加文档到集合
        self.vector_db.add_texts(
            texts=documents,
            metadatas=metadatas
        )

    def search_data(self, query: str, filter_fields: Dict[str, Any] = None, 
    k: int = 3, similarity_threshold: float = 0.6) -> List[Dict]:
        docs_with_scores = self.vector_db.similarity_search_with_score(
            query=query,
            k=k,
            filter=filter_fields
        )
        
        return [
            {
                "content": doc.page_content,
                "metadata": doc.metadata,
                "similarity": round( (score + 1) / 2, 4)  # 转换为0-1范围
            } 
            for doc, score in docs_with_scores 
            if (score + 1)/2 >= similarity_threshold  # 应用新阈值
        ]

# 使用示例
if __name__ == "__main__":
    db_manager = ChromaDBManager()
    
    # 存储示例数据
    documents = [
        """
        三年级学生学习成绩的分析报告
        第一段：三年级总体情况
        三年级共有两个班级，每班各有10名学生，总计20名学生。男女比例均衡，各有10名男生和10名女生。学生的平均年龄为8.5岁。
        第二段：学生总成绩情况
        学生的总成绩分布在239分至270分之间，平均总成绩为259.1分。各科平均分如下：数学为88.25分，英语为87分，语文为83.85分。共有10名学生的总成绩低于平均分，占比50%。
        """,
        """
        三年级学生学习成绩的分析报告
        第一部分：年级总体情况
        三年级共设2个班级（一班、二班），总人数20人，各班人数均等（每班10人）。
        1.	性别分布：
        o	男生10人（50%），女生10人（50%），男女比例约为1:1。
        o	一班：男生5人，女生5人（1:1）；二班：男生5人，女生4人（1.5:1）。
        2.	年龄特征：
        o	平均年龄：8.5岁（8岁学生10人，9岁学生10人）。
        o	年龄分布：8岁与9岁各占50%，年级结构均衡。
        第二部分：学业表现总览
        1.	总分分析：
        o	最高分：李华（270分，一班）
        o	最低分：高婷（239分，二班）
        o	年级平均分：259.1分（总分5182分/20人）
        o	分数分布：
        	高于平均分：10人（如李华、朱琳、黄杰等），占比50%
        	低于平均分：10人（如高婷、陈伟、郑美等），占比50%
        2.	学科均衡性：
        学科	语文	数学	英语
        平均分	83.1	88.7	86.4
        o	优势学科：数学（88.7分），显著高于语文与英语；
        o	薄弱学科：语文（83.1分），部分学生成绩低于80分（如高婷78分、胡勇77分）。
        """,
        """
        {{学生姓名}}学业表现分析报告
        一、总分定位
        {{学生姓名}}本次考试总分为239分，低于三年级平均分（259.1分）20.1分，位列年级末位。单科成绩对比显示：
        •	语文：78分（低于平均分5.1分，最高分92分）
        •	数学：70分（低于平均分18.7分，最高分98分）
        •	英语：91分（高于平均分4.6分，最高分94分）
        核心差距：数学成绩显著落后，与年级最高分差距达28分。
        二、短板诊断
        1.	数学学科：
        o	成绩70分，为全年级最低。
        o	其数学成绩较班级第二低分（郑美78分）仍低8分，反映系统性薄弱。
        2.	语文学科：
        o	成绩78分，低于平均分。
        o	实施路径：每日精读1篇叙事短文，完成配套习题（侧重归纳主旨与词汇填空），每周提交1篇仿写作文由教师批改。
        o	目标：期末语文成绩达到85分（接近年级平均）。
        """,
        """
        帮我生成一份关于班级学习成绩分析报告，报告要求分为2段话（200字左右，分两段）。
        第一段话，介绍一下三年级总体情况，如各班人员总数，男女比例，平均年龄等；
        第二段话，总结一下学生总成绩情况，包括但不限于学生最高分、学生最低分、平均分、学生分数与平均分的对比、各科平均分等；
        """,
        """
        作为班主任，请根据{{学生姓名}}的成绩数据，撰写一份考试成绩分析报告（200字左右，分三段）。
        要求：总分定位：对比班级平均分与单科最高分，量化差距；
        短板诊断：针对单科成绩低于班级平均分的科目，结合数据说明问题； 
        格式规范：学术报告风格，关键数据加粗，避免主观描述。
        """,

    ]
    metadatas = [
        {"category": "报告模板", "source": "班级考试成绩分析报告1"},
        {"category": "报告模板", "source": "班级考试成绩分析报告2"},
        {"category": "报告模板", "source": "学生考试成绩分析报告"},
        {"category": "提示词", "source": "班级考试成绩查询"},
        {"category": "提示词", "source": "学生考试成绩查询"},
    ]
    db_manager.store_data(documents, metadatas)
    
    # 检索示例
    # results = db_manager.search_data(
    #     query="总分定位：对比班级平均分与单科最高分，量化差距",
    #     filter_fields={"category": "提示词"},
    #     k=1,
    #     similarity_threshold=0.5
    # )
    # print("检索结果：", results)